import httpx
from openai import OpenAI

import utils


class Chat:
    def __init__(
            self,
            api_key: str,
            base_url: str | None = None,
            http_client: httpx.Client | None = None,
            model: str = "gpt-3.5-turbo",
            system_role: str = "You are a helpful assistant."
    ):
        self.conversation_list = []
        self.model = model
        self.set_system_role(system_role)
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url,
            http_client=http_client,
        )
    
    def set_system_role(self, system_role):
        self.conversation_list = [
            {"role": "system", "content": system_role},
        ]
    
    def ask_gpt(self, question: str) -> str:
        print('start ask')
        self.conversation_list.append({"role": "user", "content": question})
        completion = self.client.chat.completions.create(
            model=self.model,
            messages=self.conversation_list
        )
        print('end ask')
        ans = completion.choices[0].message.content
        self.conversation_list.append({"role": "assistant", "content": ans})
        return ans


class UML(Chat):
    
    def __init__(
            self,
            api_key: str,
            base_url: str,
            http_client: httpx.Client | None = None,
            model: str = "gpt-3.5-turbo",
            system_role: str = "You are a helpful assistant."
    ):
        self.system_role = 'You are a helpful assistant.'
        super().__init__(
            api_key=api_key,
            base_url=base_url,
            http_client=http_client,
            model=model,
            system_role=self.system_role
        )
    
    def code_to_plantuml(self, code: str) -> (str, str):
        plantuml_english = self.ask_gpt(f"{code}请为这个函数生成PlantUML语言流程图")
        plantuml_english = utils.get_puml_from_markdown(plantuml_english)
        
        ans = self.ask_gpt("请将PlantUML中的原代码替换成文字描述 使用中文")
        # print(ans)
        
        plantuml_chinese = self.ask_gpt("将这段文字描述生成PlantUML代码")
        plantuml_chinese = utils.get_puml_from_markdown(plantuml_chinese)
        return plantuml_english, plantuml_chinese
